'''
将整个数据集划分为训练集、验证集，将文件名保存到txt文件
'''
import os
import sys
import numpy as np
sys.path.append(os.path.split(sys.path[0])[0])
from config import zunYiParameter as para
from sklearn.model_selection import train_test_split

#读取文件列表
segmentation_name='segmentation'#标签
volume_name='volume'#原图
segmentation_name_list=os.listdir(para.nii_seg_path)
volume_name_list=os.listdir(para.nii_volume_path)

#需要使用替换字符串的方法，因为不同系统下，文件的默认排序是不一样的
volume_relative_path_list=[volume_name+'/'+name for name in volume_name_list]
segmentation_relative_path_list=[name.replace(volume_name,segmentation_name)
                                 for name in volume_relative_path_list]
for i in range(len(volume_relative_path_list)):
    print(volume_relative_path_list[i],segmentation_relative_path_list[i])
#原图在前，标签在后
concact_vol_seg=[vol+' '+seg  for vol,seg in zip(volume_relative_path_list,segmentation_relative_path_list)]
concact_vol_seg.sort()




#共131个数据分割占比
# train_num=100
# test_num=21
# val_num=10
# test_scale=(test_num+val_num)/(train_num+test_num+val_num)
train_list,test_list=train_test_split(concact_vol_seg,test_size =0.2,random_state = 100)#200，100,50,20
# test_list,val_list=train_test_split(test_list,test_size =(val_num/(val_num+test_num)),random_state = 100)

#保存至文件
save_path_root=os.path.dirname(para.train_nii_id_path)
if not os.path.exists(save_path_root):
    os.makedirs(save_path_root)
np.savetxt(para.train_nii_id_path,train_list,fmt = '%s')
np.savetxt(para.test_nii_id_path,test_list,fmt = '%s')
# np.savetxt(para.val_nii_id_path,val_list,fmt = '%s')
print('训练集:{}  测试集{} :'.format(len(train_list),len(test_list)))

print("保存成功,txt文件保存根路径为:"+save_path_root)
